package com.example.springai.controller;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.stabilityai.StabilityAiImageModel;
import org.springframework.ai.stabilityai.api.StabilityAiImageOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.util.Base64;

@RestController
@Slf4j
public class StabilityController {

    @Autowired
    private StabilityAiImageModel stabilityAiImageModel;

    @GetMapping(value = "/stability/generateImage")
    public void generateImage(@RequestParam(value = "message", defaultValue = "a cute cat") String message) {
        ImageResponse response = stabilityAiImageModel.call(new ImagePrompt(message, StabilityAiImageOptions.builder().withStylePreset("cinematic").withN(1).withHeight(512).withWidth(512).withResponseFormat("image/png").build()));
        convertBase64ToImage(response.getResult().getOutput().getB64Json());
    }

    public static void convertBase64ToImage(String base64String) {
        try {
            // 解码Base64字符串
            byte[] imageBytes = Base64.getDecoder().decode(base64String);

            // 将解码后的字节数组转换成图片
            ByteArrayInputStream bais = new ByteArrayInputStream(imageBytes);
            BufferedImage bufferedImage = ImageIO.read(bais);

            // 保存图片到文件系统
            ImageIO.write(bufferedImage, "png", new File("C:\\Users\\yurii\\Desktop\\output.png"));
        } catch (IOException e) {
            log.error("图片生成异常：", e);
        }
    }

}
